from core.function import Function
from core.cuda import cuda

# 矩阵的乘积
# 矩阵的计算要注意形状
class MatMul(Function):

    def forward(self, x, W):
        """
        执行矩阵乘法的前向传播。

        Args:
            x (Variable): 输入变量。
            W (ndarray): 权重参数。

        Returns:
            Variable: 前向传播的输出变量。
        """
        x, W = cuda.to_array(x), cuda.to_array(W)
        y = x.dot(W)
        return y

    def backward(self, gy):
        """
        执行矩阵乘法的反向传播。

        Args:
            gy (Variable): 相对于输出的梯度。

        Returns:
            tuple: 包含相对于输入 x, 权重 W 的梯度的元组。
        """
        x, W = self.inputs
        gx = MatMul()(gy, W.T)
        gW = MatMul()(x.T, gy)
        return gx, gW

    def matmul(self, x, W):
        """
        执行矩阵乘法的计算。

        Args:
            x (Variable): 输入变量。
            W (ndarray): 权重参数。

        Returns:
            Variable: 输出变量。
        """
        return self(x, W)
